import sys
from PyQt6.QtWidgets import QApplication, QMainWindow, QPushButton, QVBoxLayout, QHBoxLayout, QWidget, QTextEdit, QTableWidget, QTableWidgetItem, QDialog, QVBoxLayout, QHBoxLayout, QLabel, QSpinBox, QDialogButtonBox, QMessageBox, QSplitter
from PyQt6.QtCore import Qt
import numpy as np
import warnings
from Core import *
# 忽略 DeprecationWarning（可选）
warnings.filterwarnings("ignore", category=DeprecationWarning)

class Scheduler:
    def __init__(self):
        self.reservation = None
        self.forbidden_list = None
        self.collision_vector = None
        self.vector_list = None
        self.adjacency_list = None
        self.adjacency_matrix = None
        self.min_avg_weight = None
        self.best_cycle = None

    def initialize_reservation(self, reservation_data):
        """初始化预约表"""
        self.reservation = np.array(reservation_data)


class SchedulerApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Non-linear Pipeline Scheduler")
        self.setGeometry(100, 100, 1200, 500)

        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)

        # 创建主布局（垂直布局）
        main_layout = QVBoxLayout(self.central_widget)

        # 创建上半部分的水平布局（预约表 + 按钮）
        top_layout = QHBoxLayout()

        # 预约表
        self.reservation_table = QTableWidget(self)
        top_layout.addWidget(self.reservation_table, stretch=4)  # 表格占更大比例

        # 按钮布局（垂直排列）
        button_layout = QVBoxLayout()

        self.init_button = QPushButton('Initialize', self)
        self.schedule_button = QPushButton('Schedule', self)
        self.show_results_button = QPushButton('Show Results', self)

        button_layout.addWidget(self.init_button)
        button_layout.addWidget(self.schedule_button)
        button_layout.addWidget(self.show_results_button)
        button_layout.addStretch()  # 添加伸缩空间，使按钮靠上排列

        top_layout.addLayout(button_layout, stretch=1)  # 按钮占较小比例

        # 使用 QSplitter 创建可调节的分隔条
        splitter = QSplitter(Qt.Orientation.Vertical)

        # 将上半部分添加到 QSplitter 中
        top_widget = QWidget()
        top_widget.setLayout(top_layout)
        splitter.addWidget(top_widget)

        # 输出文本框
        self.output_text = QTextEdit(self)
        self.output_text.setReadOnly(True)
        splitter.addWidget(self.output_text)

        # 设置初始比例（表格部分占 70%，输出文本框占 30%）
        splitter.setSizes([700, 300])

        # 将 QSplitter 添加到主布局中
        main_layout.addWidget(splitter)

        # 连接按钮的点击事件
        self.init_button.clicked.connect(self.on_initialize)
        self.schedule_button.clicked.connect(self.on_schedule)
        self.show_results_button.clicked.connect(self.on_show_results)

        self.scheduler = Scheduler()

        # 连接 cellClicked 信号到 on_cell_clicked 方法
        self.reservation_table.cellClicked.connect(self.on_cell_clicked)  # 确保这行在所有方法定义之后

    def on_initialize(self):
        """初始化预约表"""
        # 弹出自定义对话框获取用户输入
        dialog = InitializationDialog(self)
        if dialog.exec() == QDialog.DialogCode.Accepted:
            num_function_units = dialog.num_function_units.value()
            num_time_slots = dialog.num_time_slots.value()

            # 动态设置表格的行数和列数
            self.reservation_table.setRowCount(num_function_units)
            self.reservation_table.setColumnCount(num_time_slots)

            # 生成表头标签
            header_labels = [f'T{i + 1}' for i in range(num_time_slots)]
            self.reservation_table.setHorizontalHeaderLabels(header_labels)

            # 初始化表格中的每个单元格为未打勾状态
            for i in range(num_function_units):
                for j in range(num_time_slots):
                    item = QTableWidgetItem()
                    item.setFlags(item.flags() | Qt.ItemFlag.ItemIsUserCheckable)
                    item.setCheckState(Qt.CheckState.Unchecked)
                    self.reservation_table.setItem(i, j, item)

            self.output_text.append(
                f"Reservation table initialized with {num_function_units} function units and {num_time_slots} time slots.")
        else:
            self.output_text.append("Initialization cancelled by user.")

    def update_reservation(self):
        """从表格中读取用户输入的数据并更新 Scheduler 的 reservation 属性"""
        if self.reservation_table.rowCount() == 0:
            self.output_text.append("No function units defined.")
            return

        reservation_data = []
        for i in range(self.reservation_table.rowCount()):
            row_data = []
            for j in range(self.reservation_table.columnCount()):
                item = self.reservation_table.item(i, j)
                if item.checkState() == Qt.CheckState.Checked:
                    row_data.append(1)
                else:
                    row_data.append(0)
            reservation_data.append(row_data)

        self.scheduler.initialize_reservation(reservation_data)
        self.output_text.append("Reservation table updated from user input.")

    def check_unassigned(self):
        """检查是否有未分配的功能段，以及第一个和最后一个时间段是否至少有一个功能段被分配"""
        unassigned_rows = []
        first_column_assigned = False
        last_column_assigned = False

        # 检查每一行
        for row in range(self.reservation_table.rowCount()):
            row_assigned = False
            for col in range(self.reservation_table.columnCount()):
                item = self.reservation_table.item(row, col)
                if item and item.checkState() == Qt.CheckState.Checked:
                    row_assigned = True
                    if col == 0:
                        first_column_assigned = True
                    if col == self.reservation_table.columnCount() - 1:
                        last_column_assigned = True
                    break
            if not row_assigned:
                unassigned_rows.append(row)

        # 检查最后一列是否有功能段被分配
        if not last_column_assigned:
            for row in range(self.reservation_table.rowCount()):
                item = self.reservation_table.item(row, self.reservation_table.columnCount() - 1)
                if item and item.checkState() == Qt.CheckState.Checked:
                    last_column_assigned = True
                    break

        return unassigned_rows, first_column_assigned, last_column_assigned

    def check_start_time_order(self):
        """检查功能段的起始时间段是否严格递增，不能早于或等于任何在其之前的功能段的起始时间段"""
        start_times = []  # 用于存储每个功能段的起始时间段

        # 获取每个功能段的起始时间段
        for row in range(self.reservation_table.rowCount()):
            start_time = None
            for col in range(self.reservation_table.columnCount()):
                item = self.reservation_table.item(row, col)
                if item and item.checkState() == Qt.CheckState.Checked:
                    start_time = col
                    break
            start_times.append(start_time)  # 如果该行没有打勾，则 start_time 为 None

        # 检查起始时间顺序
        time_order_violations = []
        for i in range(1, len(start_times)):
            if start_times[i] is not None:
                for j in range(i):
                    if start_times[j] is not None and start_times[i] <= start_times[j]:
                        time_order_violations.append((i, j))  # 记录违反顺序的功能段对 (当前行, 前面的行)
                        break

        return time_order_violations

    def show_unassigned_error(self, unassigned_rows, first_column_assigned, last_column_assigned,
                              time_order_violations):
        """弹出错误对话框，显示未分配的行和列以及起始时间顺序错误"""
        error_message = "以下功能段和时间段尚未进行时间分配：\n\n"

        if unassigned_rows:
            error_message += "以下功能段尚未分配时间段：\n"
            for row in unassigned_rows:
                error_message += f"  - 功能段 {row + 1}\n"

        if not first_column_assigned:
            error_message += "\n第一个时间段（T1）尚未分配任何功能段。\n"

        if not last_column_assigned:
            error_message += "\n最后一个时间段（Tn）尚未分配任何功能段。\n"

        if time_order_violations:
            error_message += "\n以下功能段的起始时间段违反了顺序规则：\n"
            for current_row, prev_row in time_order_violations:
                error_message += f"  - 功能段 {current_row + 1} 的起始时间段早于或等于功能段 {prev_row + 1} 的起始时间段。\n"

        # 弹出错误对话框
        if error_message != "以下功能段和时间段尚未进行时间分配：\n\n":
            msg_box = QMessageBox(self)
            msg_box.setIcon(QMessageBox.Icon.Warning)
            msg_box.setWindowTitle("调度错误")
            msg_box.setText(error_message)
            msg_box.exec()

    def on_schedule(self):
        """计算调度方案"""
        try:
            self.update_reservation()

            # 检查是否有未分配的功能段，以及第一个和最后一个时间段是否至少有一个功能段被分配
            unassigned_rows, first_column_assigned, last_column_assigned = self.check_unassigned()

            # 检查功能段的起始时间段是否严格递增
            time_order_violations = self.check_start_time_order()

            if unassigned_rows or not first_column_assigned or not last_column_assigned or time_order_violations:
                # 如果有未分配的功能段、指定的时间段未分配或起始时间顺序错误，弹出错误对话框并阻止调度操作
                self.show_unassigned_error(unassigned_rows, first_column_assigned, last_column_assigned,
                                           time_order_violations)
                return

            if self.scheduler.reservation is not None:
                # 生成禁止表并输出
                forbidden_list = generate_forbidden_list(self.scheduler.reservation)
                self.output_text.append(f"Forbidden List: {forbidden_list}")

                # 生成碰撞向量并输出
                collision_vector = generate_collision_vector(forbidden_list, self.scheduler.reservation.shape[1])
                self.output_text.append(f"Collision Vector: {collision_vector}")

        except Exception as e:
            self.output_text.append(f"Error during scheduling: {str(e)}")
            print(f"Error during scheduling: {str(e)}")  # 打印详细错误信息

    def on_cell_clicked(self, row, column):
        """切换单元格的打勾状态"""
        item = self.reservation_table.item(row, column)
        if item:
            current_state = item.checkState()
            new_state = Qt.CheckState.Unchecked if current_state == Qt.CheckState.Checked else Qt.CheckState.Checked
            item.setCheckState(new_state)


class InitializationDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)

        self.setWindowTitle("Initialization")

        # 创建布局
        layout = QVBoxLayout()

        # 功能段数量输入
        function_units_layout = QHBoxLayout()
        function_units_label = QLabel("Number of Function Units:")
        self.num_function_units = QSpinBox()
        self.num_function_units.setMinimum(1)
        self.num_function_units.setMaximum(100)  # 可以根据需要调整最大值
        function_units_layout.addWidget(function_units_label)
        function_units_layout.addWidget(self.num_function_units)
        layout.addLayout(function_units_layout)

        # 时间列数量输入
        time_slots_layout = QHBoxLayout()
        time_slots_label = QLabel("Number of Time Slots (T1 to Tn):")
        self.num_time_slots = QSpinBox()
        self.num_time_slots.setMinimum(1)
        self.num_time_slots.setMaximum(100)  # 可以根据需要调整最大值
        time_slots_layout.addWidget(time_slots_label)
        time_slots_layout.addWidget(self.num_time_slots)
        layout.addLayout(time_slots_layout)

        # 确定按钮
        button_box = QDialogButtonBox(QDialogButtonBox.StandardButton.Ok | QDialogButtonBox.StandardButton.Cancel)
        button_box.accepted.connect(self.accept)
        button_box.rejected.connect(self.reject)
        layout.addWidget(button_box)

        self.setLayout(layout)


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = SchedulerApp()
    window.show()
    sys.exit(app.exec())  # 使用 exec() 而不是 exec_()